package websocket

import (
	"log"
	"sync"
)

// Hub 维护活跃的客户端连接和广播消息
type Hub struct {
	// 注册的客户端
	clients map[*Client]bool

	// 从客户端注册
	register chan *Client

	// 从客户端注销
	unregister chan *Client

	// 广播消息到所有客户端
	broadcast chan []byte

	// 向特定用户发送消息
	userMessage chan *UserMessage

	// 用户ID到客户端的映射
	userClients map[string]*Client

	// 保护并发访问
	mu sync.RWMutex
}

// UserMessage 用户消息结构
type UserMessage struct {
	UserID  string
	Message []byte
}

// NewHub 创建新的Hub
func NewHub() *Hub {
	return &Hub{
		clients:     make(map[*Client]bool),
		register:    make(chan *Client),
		unregister:  make(chan *Client),
		broadcast:   make(chan []byte),
		userMessage: make(chan *UserMessage),
		userClients: make(map[string]*Client),
	}
}

// Run 启动Hub
func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			h.clients[client] = true
			if client.UserID != "" {
				h.userClients[client.UserID] = client
			}
			h.mu.Unlock()
			log.Printf("客户端已连接，用户ID: %s", client.UserID)

		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				if client.UserID != "" {
					delete(h.userClients, client.UserID)
				}
				close(client.send)
			}
			h.mu.Unlock()
			log.Printf("客户端已断开，用户ID: %s", client.UserID)

		case message := <-h.broadcast:
			h.mu.RLock()
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clients, client)
				}
			}
			h.mu.RUnlock()

		case userMsg := <-h.userMessage:
			h.mu.RLock()
			if client, ok := h.userClients[userMsg.UserID]; ok {
				select {
				case client.send <- userMsg.Message:
				default:
					close(client.send)
					delete(h.clients, client)
					delete(h.userClients, userMsg.UserID)
				}
			}
			h.mu.RUnlock()
		}
	}
}

// BroadcastToAll 广播消息给所有客户端
func (h *Hub) BroadcastToAll(message []byte) {
	h.broadcast <- message
}

// SendToUser 发送消息给特定用户
func (h *Hub) SendToUser(userID string, message []byte) {
	h.userMessage <- &UserMessage{
		UserID:  userID,
		Message: message,
	}
}

// GetConnectedUsers 获取已连接的用户列表
func (h *Hub) GetConnectedUsers() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	
	users := make([]string, 0, len(h.userClients))
	for userID := range h.userClients {
		users = append(users, userID)
	}
	return users
}
